package com.i72.processor;

import com.google.auto.service.AutoService;
import com.i72.annotation.StaticBlock;
import com.i72.utils.AnnotationUtils;
import com.sun.source.tree.Tree;
import com.sun.source.util.TreePath;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.List;

import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.Processor;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Set;

/**
 * @author Sumkor
 * @since 2021/12/30
 */
@AutoService(Processor.class)
public class StaticBlockProcessor extends BaseProcessor {

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Collections.singleton(StaticBlock.class.getCanonicalName());
    }

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        messager.printMessage(Diagnostic.Kind.NOTE, "========= StaticBlockProcessor init =========");

    }

    /**
     * 读取 StaticBlock 注解参数中的配置类，用于修改 StaticBlock 注解修饰类的语法树
     */
    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        messager.printMessage(Diagnostic.Kind.NOTE, "========= StaticBlockProcessor process =========");

        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(StaticBlock.class);
        for (Element element : elements) {
            if (element.getKind() != ElementKind.CLASS) {
                continue;
            }

            List<JCTree.JCImport> jcImports = List.nil();
            java.util.List<JCTree.JCBlock> templateBlocks = new ArrayList<>();
            TypeElement typeElement = null;


            // 获取注解中的模板类
            StaticBlock annotation = element.getAnnotation(StaticBlock.class);
            java.util.List<? extends TypeMirror> types = AnnotationUtils.getTypeMirrorFromAnnotationValue(annotation::template);
            if (types == null || types.isEmpty()) return false;
            typeElement = (TypeElement) typeUtils.asElement(types.get(0));

            // 获取模板类的语法树
            JCTree templateJcTree = null;
            if (typeElement != null) {
                templateJcTree = javacTrees.getTree(typeElement);
                // 如果获取到的语法树是空的，说明模板类已经完成编译（一般位于jar包中），是不可编辑状态了
                if (templateJcTree == null) {
                    messager.printMessage(Diagnostic.Kind.NOTE, "无法获取 " + typeElement + " 语法树");
                    return false;
                }
            }

            // 获取模板类的信息
            if (templateJcTree != null) {
                // 获取模板类的静态代码块
                templateJcTree.accept(new TreeTranslator() {
                    @Override
                    public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                        for (JCTree def : jcClassDecl.defs) {
                            if (def instanceof JCTree.JCBlock) {
                                templateBlocks.add((JCTree.JCBlock) def);
                            }
                        }
                    }
                });
                // 获取模板类的 import 信息
                TreePath treePath = trees.getPath(typeElement);
                Tree leaf = treePath.getLeaf();
                if (treePath.getCompilationUnit() instanceof JCTree.JCCompilationUnit && leaf instanceof JCTree) {
                    JCTree.JCCompilationUnit jcCompilationUnit = (JCTree.JCCompilationUnit) treePath.getCompilationUnit();
                    jcImports = jcCompilationUnit.getImports();
                }
            }

            // 修改目标类
            if (!templateBlocks.isEmpty()) {
                // 目标类添加静态代码块
                JCTree jcTree = javacTrees.getTree(element);
                jcTree.accept(new TreeTranslator() {
                    @Override
                    public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                        for (JCTree.JCBlock templateBlock : templateBlocks) {
                            jcClassDecl.defs = jcClassDecl.defs.append(templateBlock);
                        }
                        super.visitClassDef(jcClassDecl);
                    }
                });
                // 目标类添加 import 信息
                TreePath treePath = trees.getPath(element);
                Tree leaf = treePath.getLeaf();
                if (treePath.getCompilationUnit() instanceof JCTree.JCCompilationUnit && leaf instanceof JCTree) {
                    JCTree.JCCompilationUnit jcCompilationUnit = (JCTree.JCCompilationUnit) treePath.getCompilationUnit();
                    java.util.List<JCTree> jcTreeList = new ArrayList<>(jcCompilationUnit.defs);
                    for (JCTree.JCImport jcImport : jcImports) {
                        if (!jcTreeList.contains(jcImport)) {
                            jcTreeList.add(0, jcImport);
                        }
                    }
                    jcCompilationUnit.defs = List.from(jcTreeList);
                }
            }
        }
        return true;
    }

}
